#!/bin/env python
# -*- coding: utf-8 -*-
import taichi as ti
import numpy as np
import cv2
import time
import argparse
import os
import matplotlib.pyplot as plt
from tqdm import tqdm
import traceback
import sys
from taichi_dither import DitherFactory, DitherAlgorithm, BayerDitherer
import platform

import matplotlib.font_manager as fm

# 根据操作系统动态设置中文字体
system = platform.system()
if system == 'Windows':
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'sans-serif']
elif system == 'Darwin':  # macOS
    plt.rcParams['font.sans-serif'] = ['PingFang SC', 'STHeiti', 'sans-serif']
else:  # Linux 和其他系统
    plt.rcParams['font.sans-serif'] = ['AR PL UMing CN']
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class DitherTester:
    """
    抖动算法性能测试器
    测试不同算法和处理器(CPU/GPU)下的性能表现
    """
    
    def __init__(self, image_path, output_dir="test_results", fast_mode=False, use_cache=True):
        """
        初始化测试器
        
        :param image_path: 测试图像路径
        :param output_dir: 测试结果输出目录
        :param fast_mode: 快速测试模式，减少测试用例
        :param use_cache: 是否使用编译缓存，默认启用
        """
        self.image_path = image_path
        self.output_dir = output_dir
        self.fast_mode = fast_mode
        self.use_cache = use_cache
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 算法列表
        self.algorithms = [
            DitherAlgorithm.BAYER,
            DitherAlgorithm.FLOYD,
            DitherAlgorithm.ATKINSON
        ]
        
        # 色值数量列表
        self.color_counts = [2, 8] if fast_mode else [2, 4, 8, 16]
        
        # 图像大小列表
        if fast_mode:
            self.image_sizes = [
                (100, None),   # 宽度为100
                (400, None),   # 宽度为400
            ]
        else:
            self.image_sizes = [
                (100, None),   # 宽度为100
                (200, None),   # 宽度为200
                (400, None),   # 宽度为400
                (800, None),   # 宽度为800
                (1200, None),  # 宽度为1200
            ]
        
        # 结果存储
        self.results = {
            'cpu': {},
            'gpu': {}
        }
        
    def run_algorithm_test(self, use_gpu=False):
        """
        测试指定处理器下所有算法的性能
        
        :param use_gpu: 是否使用GPU
        :return: 性能结果字典
        """
        # 选择处理器
        arch = ti.gpu if use_gpu else ti.cpu
        processor = 'gpu' if use_gpu else 'cpu'
        
        print(f"\n{'='*50}")
        print(f"测试处理器: {processor.upper()}")
        print(f"{'='*50}")
        
        try:
            # 对每个算法进行测试
            for algorithm in tqdm(self.algorithms, desc="测试算法", unit="个"):
                algo_name = algorithm.value
                self.results[processor][algo_name] = {}
                
                print(f"\n测试算法: {algo_name}")
                print(f"{'-'*30}")
                
                # 创建抖动器实例（使用工厂类）
                ditherer = DitherFactory.create_ditherer(
                    algorithm=algorithm,
                    arch=arch,
                    debug=False,
                    use_cache=self.use_cache
                )
                
                # 获取系统信息
                if algo_name == self.algorithms[0].value:  # 仅在第一个算法测试时获取系统信息
                    system_info = ditherer.get_system_info()
                    print(f"系统信息:\n{system_info}\n")
                
                # 测试灰度模式
                print(f"灰度模式:")
                for color_count in tqdm(self.color_counts, desc=f"灰度模式色值", unit="个"):
                    self.results[processor][algo_name][f'gray_{color_count}'] = {}
                    
                    for width, height in tqdm(self.image_sizes, desc=f"图像尺寸", unit="个", leave=False):
                        try:
                            # 执行处理并计时
                            start_time = time.time()
                            
                            result_img = ditherer.process_image(
                                self.image_path,
                                num_colors=color_count,
                                use_color=False,
                                output_width=width,
                                output_height=height
                            )
                           
                            end_time = time.time()
                            
                            # 计算处理时间
                            process_time = end_time - start_time
                            img_size = f"{result_img.shape[1]}x{result_img.shape[0]}"
                            
                            # 存储结果
                            self.results[processor][algo_name][f'gray_{color_count}'][img_size] = process_time
                            
                            # 输出结果
                            print(f"  {color_count}色值, 尺寸 {img_size}: {process_time:.4f}秒")
                            
                            # 保存结果图像
                            output_path = os.path.join(
                                self.output_dir, 
                                f"{processor}_{algo_name}_gray_{color_count}_{img_size.replace('x', '_')}.jpg"
                            )
                            cv2.imwrite(output_path, result_img)
                        except Exception as e:
                            print(f"  处理失败: {e}")
                            traceback.print_exc(file=sys.stdout)
                
                # 测试彩色模式
                print(f"彩色模式:")
                for color_count in tqdm(self.color_counts, desc=f"彩色模式色值", unit="个"):
                    self.results[processor][algo_name][f'color_{color_count}'] = {}
                    
                    # 彩色模式只测试部分尺寸以节省时间
                    test_sizes = self.image_sizes[:2] if self.fast_mode else self.image_sizes[:3]
                    
                    for width, height in tqdm(test_sizes, desc=f"图像尺寸", unit="个", leave=False):
                        try:
                            # 执行处理并计时
                            start_time = time.time()
                        
                            result_img = ditherer.process_image(
                                self.image_path,
                                num_colors=color_count,
                                use_color=True,
                                output_width=width,
                                output_height=height
                            )
                            
                            end_time = time.time()
                            
                            # 计算处理时间
                            process_time = end_time - start_time
                            img_size = f"{result_img.shape[1]}x{result_img.shape[0]}"
                            
                            # 存储结果
                            self.results[processor][algo_name][f'color_{color_count}'][img_size] = process_time
                            
                            # 输出结果
                            print(f"  {color_count}色值, 尺寸 {img_size}: {process_time:.4f}秒")
                            
                            # 保存结果图像
                            output_path = os.path.join(
                                self.output_dir, 
                                f"{processor}_{algo_name}_color_{color_count}_{img_size.replace('x', '_')}.jpg"
                            )
                            cv2.imwrite(output_path, result_img)
                        except Exception as e:
                            print(f"  处理失败: {e}")
                            traceback.print_exc(file=sys.stdout)
            
            return self.results
        except Exception as e:
            print(f"测试处理器 {processor} 时出错: {e}")
            traceback.print_exc(file=sys.stdout)
            return self.results
    
    def run_all_tests(self):
        """
        运行所有测试
        
        :return: 所有测试结果
        """
        # 先测CPU
        self.run_algorithm_test(use_gpu=False)
        
        # 再测GPU
        try:
            self.run_algorithm_test(use_gpu=True)
        except Exception as e:
            print(f"GPU测试失败: {e}")
            print("跳过GPU测试")
            traceback.print_exc(file=sys.stdout)
        
        return self.results
    
    def generate_report(self):
        """
        生成测试报告
        
        :return: 报告文件路径
        """
        report_path = os.path.join(self.output_dir, "performance_report.txt")
        
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("抖动算法性能测试报告\n")
            f.write("=" * 50 + "\n\n")
            
            # 输出CPU测试结果
            f.write("CPU测试结果\n")
            f.write("-" * 30 + "\n")
            for algorithm in [algo.value for algo in self.algorithms]:
                if algorithm in self.results['cpu']:
                    f.write(f"\n算法: {algorithm}\n")
                    
                    # 灰度模式
                    for color_count in self.color_counts:
                        key = f'gray_{color_count}'
                        if key in self.results['cpu'][algorithm]:
                            f.write(f"  灰度{color_count}色值:\n")
                            for size, time_taken in self.results['cpu'][algorithm][key].items():
                                f.write(f"    尺寸 {size}: {time_taken:.4f}秒\n")
                    
                    # 彩色模式
                    for color_count in self.color_counts:
                        key = f'color_{color_count}'
                        if key in self.results['cpu'][algorithm]:
                            f.write(f"  彩色{color_count}色值:\n")
                            for size, time_taken in self.results['cpu'][algorithm][key].items():
                                f.write(f"    尺寸 {size}: {time_taken:.4f}秒\n")
            
            # 如果有GPU测试结果，也输出
            if self.results['gpu']:
                f.write("\n\nGPU测试结果\n")
                f.write("-" * 30 + "\n")
                for algorithm in [algo.value for algo in self.algorithms]:
                    if algorithm in self.results['gpu']:
                        f.write(f"\n算法: {algorithm}\n")
                        
                        # 灰度模式
                        for color_count in self.color_counts:
                            key = f'gray_{color_count}'
                            if key in self.results['gpu'][algorithm]:
                                f.write(f"  灰度{color_count}色值:\n")
                                for size, time_taken in self.results['gpu'][algorithm][key].items():
                                    f.write(f"    尺寸 {size}: {time_taken:.4f}秒\n")
                        
                        # 彩色模式
                        for color_count in self.color_counts:
                            key = f'color_{color_count}'
                            if key in self.results['gpu'][algorithm]:
                                f.write(f"  彩色{color_count}色值:\n")
                                for size, time_taken in self.results['gpu'][algorithm][key].items():
                                    f.write(f"    尺寸 {size}: {time_taken:.4f}秒\n")
            
            # CPU vs GPU 性能比较
            if self.results['gpu']:
                f.write("\n\nCPU vs GPU 性能比较\n")
                f.write("-" * 30 + "\n")
                
                for algorithm in [algo.value for algo in self.algorithms]:
                    if algorithm in self.results['gpu']:
                        f.write(f"\n算法: {algorithm}\n")
                        
                        # 灰度模式比较
                        for color_count in self.color_counts:
                            key = f'gray_{color_count}'
                            if key in self.results['cpu'][algorithm] and key in self.results['gpu'][algorithm]:
                                f.write(f"  灰度{color_count}色值:\n")
                                for size in self.results['cpu'][algorithm][key]:
                                    if size in self.results['gpu'][algorithm][key]:
                                        cpu_time = self.results['cpu'][algorithm][key][size]
                                        gpu_time = self.results['gpu'][algorithm][key][size]
                                        speedup = cpu_time / gpu_time if gpu_time > 0 else 0
                                        f.write(f"    尺寸 {size}: 加速比 {speedup:.2f}x (CPU: {cpu_time:.4f}秒, GPU: {gpu_time:.4f}秒)\n")
                        
                        # 彩色模式比较
                        for color_count in self.color_counts:
                            key = f'color_{color_count}'
                            if key in self.results['cpu'][algorithm] and key in self.results['gpu'][algorithm]:
                                f.write(f"  彩色{color_count}色值:\n")
                                for size in self.results['cpu'][algorithm][key]:
                                    if size in self.results['gpu'][algorithm][key]:
                                        cpu_time = self.results['cpu'][algorithm][key][size]
                                        gpu_time = self.results['gpu'][algorithm][key][size]
                                        speedup = cpu_time / gpu_time if gpu_time > 0 else 0
                                        f.write(f"    尺寸 {size}: 加速比 {speedup:.2f}x (CPU: {cpu_time:.4f}秒, GPU: {gpu_time:.4f}秒)\n")
        
        print(f"\n性能测试报告已生成: {report_path}")
        return report_path
    
    def generate_charts(self):
        """
        生成性能图表
        
        :return: 图表文件路径列表
        """
        chart_files = []
        
        # 为每个算法生成性能图表
        for algorithm in [algo.value for algo in self.algorithms]:
            if algorithm not in self.results['cpu']:
                continue
                
            # 灰度模式
            plt.figure(figsize=(12, 8))
            plt.title(f"{algorithm.capitalize()} 算法 - 灰度模式性能")
            plt.xlabel("图像尺寸")
            plt.ylabel("处理时间 (秒)")
            
            for color_count in self.color_counts:
                key = f'gray_{color_count}'
                if key in self.results['cpu'][algorithm]:
                    sizes = list(self.results['cpu'][algorithm][key].keys())
                    cpu_times = [self.results['cpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, cpu_times, 'o-', label=f"CPU {color_count}色值")
                
                if 'gpu' in self.results and algorithm in self.results['gpu'] and key in self.results['gpu'][algorithm]:
                    sizes = list(self.results['gpu'][algorithm][key].keys())
                    gpu_times = [self.results['gpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, gpu_times, 'x--', label=f"GPU {color_count}色值")
            
            plt.legend()
            plt.grid(True)
            plt.xticks(rotation=45)
            
            # 保存图表
            chart_path = os.path.join(self.output_dir, f"{algorithm}_gray_performance.png")
            plt.savefig(chart_path, dpi=100, bbox_inches='tight')
            chart_files.append(chart_path)
            plt.close()
            
            # 彩色模式
            plt.figure(figsize=(12, 8))
            plt.title(f"{algorithm.capitalize()} 算法 - 彩色模式性能")
            plt.xlabel("图像尺寸")
            plt.ylabel("处理时间 (秒)")
            
            for color_count in self.color_counts:
                key = f'color_{color_count}'
                if key in self.results['cpu'][algorithm]:
                    sizes = list(self.results['cpu'][algorithm][key].keys())
                    cpu_times = [self.results['cpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, cpu_times, 'o-', label=f"CPU {color_count}色值")
                
                if 'gpu' in self.results and algorithm in self.results['gpu'] and key in self.results['gpu'][algorithm]:
                    sizes = list(self.results['gpu'][algorithm][key].keys())
                    gpu_times = [self.results['gpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, gpu_times, 'x--', label=f"GPU {color_count}色值")
            
            plt.legend()
            plt.grid(True)
            plt.xticks(rotation=45)
            
            # 保存图表
            chart_path = os.path.join(self.output_dir, f"{algorithm}_color_performance.png")
            plt.savefig(chart_path, dpi=100, bbox_inches='tight')
            chart_files.append(chart_path)
            plt.close()
        
        # 算法间的比较图
        # CPU灰度模式
        plt.figure(figsize=(12, 8))
        plt.title("算法比较 - CPU灰度模式 (8色值)")
        plt.xlabel("图像尺寸")
        plt.ylabel("处理时间 (秒)")
        
        for algorithm in [algo.value for algo in self.algorithms]:
            if algorithm not in self.results['cpu']:
                continue
                
            key = 'gray_8'
            if key in self.results['cpu'][algorithm]:
                sizes = list(self.results['cpu'][algorithm][key].keys())
                times = [self.results['cpu'][algorithm][key][size] for size in sizes]
                plt.plot(sizes, times, 'o-', label=f"{algorithm.capitalize()}")
        
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        
        # 保存图表
        chart_path = os.path.join(self.output_dir, "cpu_gray_algorithm_comparison.png")
        plt.savefig(chart_path, dpi=100, bbox_inches='tight')
        chart_files.append(chart_path)
        plt.close()
        
        # CPU彩色模式
        plt.figure(figsize=(12, 8))
        plt.title("算法比较 - CPU彩色模式 (8色值)")
        plt.xlabel("图像尺寸")
        plt.ylabel("处理时间 (秒)")
        
        for algorithm in [algo.value for algo in self.algorithms]:
            if algorithm not in self.results['cpu']:
                continue
                
            key = 'color_8'
            if key in self.results['cpu'][algorithm]:
                sizes = list(self.results['cpu'][algorithm][key].keys())
                times = [self.results['cpu'][algorithm][key][size] for size in sizes]
                plt.plot(sizes, times, 'o-', label=f"{algorithm.capitalize()}")
        
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        
        # 保存图表
        chart_path = os.path.join(self.output_dir, "cpu_color_algorithm_comparison.png")
        plt.savefig(chart_path, dpi=100, bbox_inches='tight')
        chart_files.append(chart_path)
        plt.close()
        
        # 如果有GPU结果，则生成GPU算法比较图
        if self.results['gpu']:
            # GPU灰度模式
            plt.figure(figsize=(12, 8))
            plt.title("算法比较 - GPU灰度模式 (8色值)")
            plt.xlabel("图像尺寸")
            plt.ylabel("处理时间 (秒)")
            
            for algorithm in [algo.value for algo in self.algorithms]:
                if algorithm not in self.results['gpu']:
                    continue
                    
                key = 'gray_8'
                if key in self.results['gpu'][algorithm]:
                    sizes = list(self.results['gpu'][algorithm][key].keys())
                    times = [self.results['gpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, times, 'o-', label=f"{algorithm.capitalize()}")
            
            plt.legend()
            plt.grid(True)
            plt.xticks(rotation=45)
            
            # 保存图表
            chart_path = os.path.join(self.output_dir, "gpu_gray_algorithm_comparison.png")
            plt.savefig(chart_path, dpi=100, bbox_inches='tight')
            chart_files.append(chart_path)
            plt.close()
            
            # GPU彩色模式
            plt.figure(figsize=(12, 8))
            plt.title("算法比较 - GPU彩色模式 (8色值)")
            plt.xlabel("图像尺寸")
            plt.ylabel("处理时间 (秒)")
            
            for algorithm in [algo.value for algo in self.algorithms]:
                if algorithm not in self.results['gpu']:
                    continue
                    
                key = 'color_8'
                if key in self.results['gpu'][algorithm]:
                    sizes = list(self.results['gpu'][algorithm][key].keys())
                    times = [self.results['gpu'][algorithm][key][size] for size in sizes]
                    plt.plot(sizes, times, 'o-', label=f"{algorithm.capitalize()}")
            
            plt.legend()
            plt.grid(True)
            plt.xticks(rotation=45)
            
            # 保存图表
            chart_path = os.path.join(self.output_dir, "gpu_color_algorithm_comparison.png")
            plt.savefig(chart_path, dpi=100, bbox_inches='tight')
            chart_files.append(chart_path)
            plt.close()
            
            # CPU vs GPU 加速比对比图
            plt.figure(figsize=(12, 8))
            plt.title("CPU vs GPU 加速比对比 (彩色8色值)")
            plt.xlabel("图像尺寸")
            plt.ylabel("加速比 (CPU时间/GPU时间)")
            
            for algorithm in [algo.value for algo in self.algorithms]:
                if algorithm not in self.results['gpu']:
                    continue
                    
                key = 'color_8'
                speedups = []
                sizes = []
                
                if key in self.results['cpu'][algorithm] and key in self.results['gpu'][algorithm]:
                    for size in self.results['cpu'][algorithm][key]:
                        if size in self.results['gpu'][algorithm][key]:
                            cpu_time = self.results['cpu'][algorithm][key][size]
                            gpu_time = self.results['gpu'][algorithm][key][size]
                            if gpu_time > 0:
                                speedup = cpu_time / gpu_time
                                speedups.append(speedup)
                                sizes.append(size)
                
                if speedups:
                    plt.plot(sizes, speedups, 'o-', label=f"{algorithm.capitalize()}")
            
            plt.axhline(y=1.0, color='r', linestyle='--', label="CPU=GPU")
            plt.legend()
            plt.grid(True)
            plt.xticks(rotation=45)
            
            # 保存图表
            chart_path = os.path.join(self.output_dir, "gpu_speedup_comparison.png")
            plt.savefig(chart_path, dpi=100, bbox_inches='tight')
            chart_files.append(chart_path)
            plt.close()
        
        print(f"\n已生成 {len(chart_files)} 个性能图表")
        return chart_files

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='测试抖动算法的性能')
    parser.add_argument('--image', type=str, default='test.jpg', help='测试图像路径')
    parser.add_argument('--output', type=str, default='test_results', help='输出结果目录')
    parser.add_argument('--fast', action='store_true', help='快速测试模式，减少测试用例')
    parser.add_argument('--no-cache', action='store_true', help='禁用Taichi编译缓存')
    args = parser.parse_args()
    
    # 创建测试器
    tester = DitherTester(
        args.image, 
        args.output, 
        fast_mode=args.fast,
        use_cache=not args.no_cache
    )
    
    # 运行测试
    print("开始性能测试...")
    tester.run_all_tests()
    
    # 生成报告
    tester.generate_report()
    
    # 生成图表
    tester.generate_charts()
    
    print("\n测试完成！")

if __name__ == "__main__":
    main() 